package org.zetaframework.mq;

import cn.dev33.satoken.stp.StpUtil;
import org.springframework.cloud.stream.function.StreamBridge;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;
import org.zetaframework.core.utils.ContextUtil;

import java.util.function.Consumer;

/**
 * Mq消息发送者
 * <p>
 * 说明:<br>
 * StreamBridge 发消息的简单封装
 *
 * @author gcc
 */
@Component
public class MessageSender {

    public static final String ROUTING_KEY = "routingKey";

    private final StreamBridge streamBridge;
    public MessageSender(StreamBridge streamBridge) {
        this.streamBridge = streamBridge;
    }


    /**
     * 创建消息构建器，指定目标
     *
     * @param bindingName 目标绑定名称
     * @return 消息构建器
     */
    public MessageBuilderWrapper to(String bindingName) {
        return new MessageBuilderWrapper(bindingName);
    }


    /**
     * 消息构建器包装类
     */
    public class MessageBuilderWrapper {

        /** 绑定名称 */
        private final String bindingName;

        /** 路由key */
        private String routingKey = "#";

        public MessageBuilderWrapper(String bindingName) {
            this.bindingName = bindingName;
        }

        /**
         * 设置路由key
         * <p>
         * 说明：<br>
         * 不主动设置，则默认发送给所有消费者
         *
         * @param routingKey 路由key
         * @return MessageBuilderWrapper
         */
        public MessageBuilderWrapper routingKey(String routingKey) {
            this.routingKey = routingKey;
            return this;
        }

        /**
         * 发送消息 给指定消费者
         *
         * @param message 消息数据
         * @param <T> 消息类型
         */
        public <T> void sendMessage(T message) {
            sendMessage(message, null);
        }


        /**
         * 发送消息 给指定消费者
         *
         * @param message 消息数据
         * @param headerBuilder 自定义消息头
         * @param <T> 消息类型
         */
        public <T> void sendMessage(T message, Consumer<MessageBuilder<T>> headerBuilder) {
            // 创建消息
            MessageBuilder<T> msg = MessageBuilder.withPayload(message);
            // 携带token
            msg.setHeader(StpUtil.getTokenName(), ContextUtil.getToken());
            // 添加路由key
            msg.setHeader(ROUTING_KEY, this.routingKey);
            // 添加自定义header
            if (headerBuilder != null) {
                headerBuilder.accept(msg);
            }

            // 发送消息
            streamBridge.send(this.bindingName, msg.build());
        }

        /**
         * 异步发送消息 给指定消费者
         *
         * @param message 消息数据
         * @param <T> 消息类型
         */
        @Async
        public <T> void sendMessageAsync(T message) {
            sendMessageAsync(message, null);
        }

        /**
         * 异步发送消息 给指定消费者
         *
         * @param message 消息数据
         * @param headerBuilder 自定义消息头
         * @param <T> 消息类型
         */
        @Async
        public <T> void sendMessageAsync(T message, Consumer<MessageBuilder<T>> headerBuilder) {
            sendMessage(message, headerBuilder);
        }
    }

}
